package com.codejiwei.flink.broadcast;

import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.java.typeutils.MapTypeInfo;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.util.Map;

/**
 * author: codejiwei
 * date: 2023/8/2
 * desc: broadcast demo
 **/
public class Flink_BroadCast_Test {
    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        DataStreamSource<String> source = env.socketTextStream("localhost", 8889);

        MapStateDescriptor<String, Map<String, String>> broadcastState = new MapStateDescriptor<>(
                "RulesBroadcastState", BasicTypeInfo.STRING_TYPE_INFO, new MapTypeInfo<String, String>(
                BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO
        ));

        BroadcastStream<Map<String, String>> ruleDS = env.addSource(new RuleBroadcastFetcher())
                .broadcast(broadcastState);


        source.connect(ruleDS)
                .process(new BroadcastProcessFunction<String, Map<String, String>, String>() {
                    @Override
                    public void processElement(String s, ReadOnlyContext readOnlyContext, Collector<String> collector) throws Exception {

                        Map<String, String> rule_broadcast = readOnlyContext.getBroadcastState(broadcastState).get("rule_broadcast");

                        for (Map.Entry<String, String> entry : rule_broadcast.entrySet()) {
                            if (s.equals(entry.getValue())) {
                                collector.collect(entry.getKey() + "---" + entry.getValue() + "---" + s);
                            }
                        }

                    }

                    @Override
                    public void processBroadcastElement(Map<String, String> stringStringMap, Context context, Collector<String> collector) throws Exception {
                        context.getBroadcastState(broadcastState).put("rule_broadcast", stringStringMap);
                    }
                }).print();


//        source.print();

        env.execute();


    }
}
